"""
authenticators: the server instance accepts an authenticator object,
which is basically any callable (i.e., a function) that takes the newly
connected socket and "authenticates" it. 

the authenticator should return a socket-like object with its associated 
credentials (a tuple), or raise AuthenticationError if it fails.

a very trivial authenticator might be

    def magic_word_authenticator(sock):
        if sock.recv(5) != "Ma6ik":
            raise AuthenticationError("wrong magic word")
        return sock, None
    
    s = ThreadedServer(...., authenticator = magic_word_authenticator)

your authenticator can return any socket-like object. for instance, it may 
authenticate the client and return a TLS/SSL-wrapped socket object that 
encrypts the transport.

the credentials returned alongside with the new socket can be any object.
it will be stored in the rpyc connection configruation under the key
"credentials", and may be used later by the service logic. if no credentials
are applicable, just return None as in the example above.

rpyc includes integration with PyCrypto [http://www.amk.ca/python/code/crypto.html],
and provides two common authenticators out-of-the-box:
 * SSL (encrypted, but not authenticated)
 * Username-password authenticator over SSL
 * Public-key authenticator over SSL
"""
import os
from rpyc.utils import security


class AuthenticationError(Exception):
    pass

class BaseAuthenticator(object):
    @classmethod
    def get_name(cls):
        raise NotImplementedError()
    def handshake_server(self, sock):
        raise NotImplementedError()
    def handshake_client(self, sock, **kwargs):
        raise NotImplementedError()


class NullAuthenticator(BaseAuthenticator):
    @classmethod
    def get_name(cls):
        raise "null"
    def handshake_server(self, sock):
        return sock, None
    def handshake_client(self, sock):
        return sock, None


class TlsEncrypter(BaseAuthenticator):
    """creates a TLS-encrypted connection, without authentication"""
    
    def __init__(self, certfile = None, cafile = None, keyfile = None):
        self.certfile = certfile
        self.cafile = cafile
        self.keyfile = keyfile
    @classmethod
    def get_name(cls):
        raise "tls"
    
    def handshake_server(self, sock):
        return security.TlsProtocol.handshake_server(sock), None
    
    def handshake_client(self, sock):
        return security.TlsProtocol.handshake_client(sock)


class UserDBAuthenticator(BaseAuthenticator):
    def __init__(self, userdb):
        self.userdb = userdb
    @classmethod
    def get_name(cls):
        return "tls+user"
    @classmethod
    def from_file(cls, filename):
        return cls(security.UserDB.load(filename))
    @classmethod
    def from_dict(cls, users):
        udb = security.UserDB({})
        for un, pw in users.iteritems():
            udb.set_password(un, pw)
        return cls(udb)
    
    def handshake_server(self, sock):
        try:
            sock2 = security.TlsProtocol.handshake_server(sock)
            sock3, username = security.UserProtocol.handshake_server(sock2, self.userdb)
        except Exception, ex:
            raise AuthenticationError(str(ex))
        return sock3, username
    
    def handshake_client(self, sock, username, password):
        sock2 = security.TlsProtocol.handshake_client(sock)
        return security.UserProtocol.handshake_client(sock2, username, password)


class PubkeyDBAuthenticator(BaseAuthenticator):
    def __init__(self, pubkeydb):
        self.pubkeydb = pubkeydb
    @classmethod
    def get_name(cls):
        return "tls+pubkey"
    @classmethod
    def from_file(cls, filename):
        return cls(security.PubkeyDB.load(filename))
    
    def handshake_server(self, sock):
        try:
            sock2 = security.TlsProtocol.handshake_server(sock)
            sock3, pubkey = security.PubkeyProtocol.handshake_server(sock2, self.pubkeydb)
        except Exception, ex:
            raise AuthenticationError(str(ex))
        return sock3, pubkey

    def handshake_client(self, sock, host, allow_unknown_host = False):
        if allow_unknown_host:
            default_key = security.AnyRSAKey
        else:
            default_key = security.BadRSAKey
        host_key = self.pubkeydb.get(host, default_key)
        sock2 = security.TlsProtocol.handshake_client(sock)
        return security.PubkeyProtocol.handshake_client(sock2, 
            self.pubkeydb.private_key, host_key)




